from torch import nn


class SimpleNet(nn.Module):

    def __init__(self, num_classes, in_channels=3):
        super().__init__()

        self.num_classes = num_classes
        self.extractor = nn.Sequential(
            nn.Conv2d(in_channels, 128, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(128, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.AdaptiveAvgPool2d(1)
        )
        self.classifier = nn.Sequential(
            nn.Linear(256, 128),
            nn.ReLU(inplace=True),
            nn.Linear(128, num_classes)
        )

    def forward(self, x):
        feat = self.extractor(x)
        # (batch, channels)
        feat = feat.view(feat.shape[0], feat.shape[1])
        return self.classifier(feat)
